/*
 * Project CHEN Rendering Engine : /bsdf.hpp
 * Created: 2021-02-19, Last modified:  2021-02-19
 * 
 * This is the bsdf header file.
 * 
 * Copyright (C) 2021 Hineven (hineven@pku.edu.cn)
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#ifndef _H_BSDF
#define _H_BSDF
#include "heads.hpp"
#include "spectrum.hpp"
#include "optics.hpp"
#include "buffer.hpp"
namespace chen {
// Bidirectional Transmision Distribution Function
// & 
// Bidirectional Reflection Distribution Function
class BxDF {
public:
	unsigned flags;
	inline BxDF (int flags): flags(flags) {}
    virtual Spectrum f (const Vector3 & wi, const Vector3 & wo) const = 0;
    virtual Float pdf (const Vector3 & wi, const Vector3 & wo) const = 0;
    virtual Spectrum sample (
        const Vector3 & wx, const Point2 & u,
        Vector3 * wy, Float * pdf, bool dir, unsigned * tags
    ) const = 0;
    virtual BxDF * clone () const = 0;
    inline virtual ~BxDF () {}
} ;

// Bidirectional Scattering Distribution Function
class BSDF {
public:
// Maximum number of integrated bxdfs
constexpr static int MAX_COMPONENTS = 4;
    // Surface shading normal, shading tangent space y axis, and geometry normal
    Normal ng, ns, nsy;
    // Relative refraction index
    // Inside refraction index / Outside (ng side) refraction index
    // This value is 1 if the material of the hit point is opaque
    const Float eta; 
	BxDF * bxdfs[MAX_COMPONENTS];
	int ncomponents;
	inline Spectrum f (const Vector3 & wi, const Vector3 & wo) const {
        bool reflect = ng.sideWith(wi)^ng.sideWith(wo); // resolve light leakings
        Vector3 twi = ns.coord(nsy, wi);
        Vector3 two = ns.coord(nsy, wo);
		Spectrum ret = 0;
		for(int i = 0; i<ncomponents; i++) 
            if((hasFlag(bxdfs[i]->flags, FBSDF_REFLECTION) && reflect)
            || (hasFlag(bxdfs[i]->flags, FBSDF_TRANSMISSION && !reflect)))
			    ret += bxdfs[i]->f(twi, two);
		return ret;
	}
	inline Float pdf (const Vector3 & wi, const Vector3 & wo) const {
        Vector3 twi = ns.coord(nsy, wi);
        Vector3 two = ns.coord(nsy, wo);
		Float ret = 0;
        // No need to check light leaking & black spots because the sample
        // function itself samples directions according to this rule.
		for(int i = 0; i<ncomponents; i++)
			ret += bxdfs[i]->pdf(twi, two);
		return ncomponents ? ret/ncomponents : 0;
	}
    // Sample values & directions of the bsdf function f(wx, wy), where wx is the
    // incident light direction and wy is the outgoing light direction.
    // if _dir_ is true, we use wy to sample wx instead.
	inline Spectrum sample (
        const Vector3 & wx, const Point2 & u,
        Vector3 * wy, Float * pdf, bool dir, unsigned * tags
    ) const {
        Vector3 twx = ns.coord(nsy, wx);
        int p = min((int)(u.x*ncomponents), ncomponents-1);
		Point2 nu(
			min((u.x-p)*ncomponents, float_one_minus),
			u.y
		);
        Float tpdf;
        Vector3 twy;
        Spectrum ret = bxdfs[p]->sample(twx, nu, &twy, &tpdf, dir, tags);
        *wy = ns.rcoord(nsy, twy);
        // Special processing for specular surface.
        // More concentrated processings
        if(!hasFlag(bxdfs[p]->flags, FBSDF_SPECULAR)) { 
            for(int i = 0; i<ncomponents; i++)
                if(i != p) {
                    if(!dir) tpdf += bxdfs[i]->pdf(twx, twy);
                    else tpdf += bxdfs[i]->pdf(twy, twx);
                }
            if(ncomponents != 1) tpdf /= ncomponents;
        }
        *pdf = tpdf;
        // Light leaking / black spots are not considered when calculating
        // specular surfaces. ???
        if(!hasFlag(bxdfs[p]->flags, FBSDF_SPECULAR)) {
            ret = 0;
            bool reflect = ng.sideWith(wx)^ng.sideWith(*wy);
            for(int i = 0; i<ncomponents; i++) {
                if((hasFlag(bxdfs[i]->flags, FBSDF_REFLECTION) && reflect)
                || (hasFlag(bxdfs[i]->flags, FBSDF_TRANSMISSION && !reflect))) {
                    if(!dir) ret += bxdfs[i]->f(twx, twy);
                    else ret += bxdfs[i]->f(twy, twx);
                }
            }
        }
        return ret;
	}
    // The normal here should be shading normal
	inline BSDF (const Normal & ng, const Normal & ns, const Normal & nsy, const Float & eta):
        ng(ng), ns(ns), nsy(nsy), eta(eta), ncomponents(0) {
	}
	inline void append (BxDF * bxdf) {
		assert(ncomponents != MAX_COMPONENTS-1);
		bxdfs[ncomponents++] = bxdf;
	}
	inline ~BSDF () {
		for(int i = 0; i<ncomponents; i++)
			bfree(bxdfs[i]);
	}
    inline BSDF * clone () const {
        BSDF * ret = balloc(BSDF)(ng, ns, nsy, eta);
        for(int i = 0; i<ncomponents; i++)
            ret->bxdfs[i] = bxdfs[i]->clone();
        ret->ncomponents = ncomponents;
        return ret;
    }
    inline int count (unsigned flags) const {
        switch(ncomponents) {
            case 0: return 0;
            case 1: return (bool)((bxdfs[0]->flags)&flags);
            case 2: return (((bxdfs[0]->flags)&flags) ? 1 : 0)
                         + (((bxdfs[1]->flags)&flags) ? 1 : 0);
            case 3: return (((bxdfs[0]->flags)&flags) ? 1 : 0)
                         + (((bxdfs[1]->flags)&flags) ? 1 : 0)
                         + (((bxdfs[2]->flags)&flags) ? 1 : 0);
            case 4: return (((bxdfs[0]->flags)&flags) ? 1 : 0)
                         + (((bxdfs[1]->flags)&flags) ? 1 : 0)
                         + (((bxdfs[2]->flags)&flags) ? 1 : 0)
                         + (((bxdfs[3]->flags)&flags) ? 1 : 0);
            default: throw "im fucked";
        }
    }
};
class BSSRDF {
    // TODO not implemented
};


class DiffuseBxDF: public BxDF {
public:
    Spectrum identity;
    inline DiffuseBxDF (const Spectrum & identity):
        BxDF::BxDF(FBSDF_REFLECTION), identity(identity) {}
    inline Spectrum f (const Vector3 & wi, const Vector3 & wo) const {
        if(wi.z*wo.z < 0) return identity/(Float)(2*M_PI);
        return 0;
    }
    inline Float pdf(const Vector3 & wi, const Vector3 & wo) const 
    {
        if(wi.z*wo.z < 0) return 1/(Float)(2*M_PI);
        return 0;
    }
    inline Spectrum sample (
        const Vector3 & wx, const Point2 & u,
        Vector3 * wy, Float * pdf, bool dir, unsigned * tags
    ) const {
        // TODO optmize cosine weighted sample
        unused(dir);
        if(tags != nullptr) *tags = flags;
        Float phi = (Float)(2.0*M_PI)*u.y;
        Float nz = u.x, nr = std::sqrt(std::max(0.0, 1.0-nz*nz));
        Float nx = nr*cos(phi);
        Float ny = nr*sin(phi);
        wy->x = nx, wy->y = ny, wy->z = (wx.z>0) ? -nz : nz;
        *pdf = 1/(Float)(2*M_PI);
        return identity/(Float)(2*M_PI);
    }
    inline BxDF * clone () const {
        return balloc(DiffuseBxDF)(identity);
    }
};



class FresnelBxDF: public BxDF {
public:
    Spectrum identity;
    Float glassEta;
    Fresnel fresnel;
    FresnelBxDF (const Spectrum & identity, const Fresnel & fr):
		BxDF::BxDF(FBSDF_DELTA|FBSDF_TRANSMISSION|FBSDF_REFLECTION|FBSDF_SPECULAR),
        identity(identity), fresnel(fr) {}
    inline Spectrum f (const Vector3 & wi, const Vector3 & wo) const {
        unused(wi);
        unused(wo);
        return 0;
    }
    inline Float pdf (const Vector3 & wi, const Vector3 & wo) const {
        unused(wi);
        unused(wo);
        return 0;
    }
    inline Spectrum sample (
        const Vector3 & wx, const Point2 & u,
        Vector3 * wy, Float * pdf, bool dir, unsigned * tags
    ) const {
        unused(dir); // Refraction is symmetric
        Float cosIn = 0, cosTran = wx.z;
        Float sinIn, sinTran = std::sqrt(max(0.0, 1.0-cosTran*cosTran));
        if(wx.z < 0) {
            sinIn = fresnel.etaNegative*sinTran/fresnel.etaPositive;
        } else {
            sinIn = fresnel.etaPositive*sinTran/fresnel.etaNegative;
        }
        Float reflectance;
        if(sinIn >= 1) reflectance = 1;
        else {
            cosIn = std::sqrt(max(0.0, 1.0-sinIn*sinIn));
            reflectance = fresnel.evaluate(cosIn, abs(cosTran), cosTran < 0);
        }
        if(u.x > reflectance) { // Transmission
            if(tags != nullptr)
                *tags = FBSDF_DELTA|FBSDF_SPECULAR|FBSDF_TRANSMISSION;
            // For better performance, we compute transmit direction inline
            // *wy = fresnel.transmitDirection(wx);
            wy->z = wx.z < 0 ? -cosIn : cosIn;
            Float s = std::sqrt(1-wy->z*wy->z)/std::sqrt(wx.x*wx.x+wx.y*wx.y);
            wy->x = wx.x*s, wy->y = wx.y*s;
            *pdf = 1.0-reflectance;
            Float fac;
            if(wx.z < 0) fac =
                (fresnel.etaNegative*fresnel.etaNegative)
              / (fresnel.etaPositive*fresnel.etaPositive);
            else fac =
                (fresnel.etaPositive*fresnel.etaPositive)
              / (fresnel.etaNegative*fresnel.etaNegative);
            return identity*((1.0-reflectance)*fac/abs(wy->z));
        } else { // Reflection
            if(tags != nullptr)
                *tags = FBSDF_DELTA|FBSDF_SPECULAR|FBSDF_REFLECTION;
            wy->z = -wx.z;
            wy->x = wx.x;
            wy->y = wx.y;
            *pdf = reflectance;
            return identity*(reflectance/abs(wy->z));
        }
    }
    inline BxDF * clone () const {
        return balloc(FresnelBxDF)(identity, fresnel);
    }
};
}
#endif